# -*- coding: utf-8 -*-
"""
员工流失预测 —— 特征工程模块（增强版 + 可视化）
- 无标签泄露
- 训练/测试一致处理
- 增强特征表达力
- 更稳健的特征筛选
- 新增：特征 AUC 可视化
- 【已移除所有日志输出】
"""

import pandas as pd
import numpy as np
from sklearn.metrics import roc_auc_score
import matplotlib.pyplot as plt
import os

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def engineer_features(df, is_train=False):
    """
    对输入 DataFrame 进行业务特征构造（无标签依赖）
    """
    X = df.copy()

    # === 1. 缺失值填充 ===
    for col in X.columns:
        if X[col].dtype == 'object':
            mode_val = X[col].mode()
            fill_val = mode_val[0] if not mode_val.empty else 'Missing'
            X[col] = X[col].fillna(fill_val)
        else:
            X[col] = X[col].fillna(X[col].mean())

    original_cols = set(X.columns)

    # === 2. 数值特征构造（安全、无未来信息）===
    if 'Age' in X.columns:
        X['AgeGroup'] = pd.cut(X['Age'], bins=[0, 25, 35, 50, 100], labels=[0, 1, 2, 3]).astype(int)

    if 'MonthlyIncome' in X and 'TotalWorkingYears' in X:
        X['IncomePerYearWorked'] = X['MonthlyIncome'] / (X['TotalWorkingYears'] + 1)

    if 'YearsAtCompany' in X and 'TotalWorkingYears' in X:
        X['LoyaltyRatio'] = X['YearsAtCompany'] / (X['TotalWorkingYears'] + 1e-6)

    if 'YearsAtCompany' in X and 'JobLevel' in X:
        X['PromotionSpeed'] = X['YearsAtCompany'] / (X['JobLevel'] + 1e-6)

    if 'OverTime' in X and 'MonthlyIncome' in X:
        ot_map = {'Yes': 1, 'No': 0}
        X['OverTimeNum'] = X['OverTime'].map(ot_map).fillna(0)
        X['OverTimeStress'] = X['OverTimeNum'] * X['MonthlyIncome']

    if 'DistanceFromHome' in X and 'WorkLifeBalance' in X:
        X['CommuteStress'] = X['DistanceFromHome'] * (4 - X['WorkLifeBalance'])

    if 'JobSatisfaction' in X and 'EnvironmentSatisfaction' in X:
        X['OverallSatisfaction'] = (X['JobSatisfaction'] + X['EnvironmentSatisfaction']) / 2.0

    if 'NumCompaniesWorked' in X and 'TotalWorkingYears' in X:
        X['JobHoppingRate'] = X['NumCompaniesWorked'] / (X['TotalWorkingYears'] + 1e-6)

    if 'StockOptionLevel' in X and 'MonthlyIncome' in X:
        X['WealthSignal'] = X['StockOptionLevel'] * X['MonthlyIncome']

    # === 3. 类别变量增强：频率编码（无 y 依赖）===
    cat_cols = X.select_dtypes(include=['object']).columns.tolist()
    for col in cat_cols:
        freq_map = X[col].value_counts(normalize=True).to_dict()
        X[f'{col}_FreqEnc'] = X[col].map(freq_map)

    return X


def _compute_feature_auc_scores(X, y):
    """计算每个特征的单变量 AUC（或 1-AUC，取更远离 0.5 的方向）"""
    auc_scores = {}
    for col in X.columns:
        if X[col].nunique() < 2:
            auc_scores[col] = 0.5
            continue
        try:
            if X[col].dtype == 'object':
                temp = X[col].astype('category').cat.codes
            else:
                temp = X[col]
            auc = roc_auc_score(y, temp)
            auc_scores[col] = max(auc, 1 - auc)
        except Exception:
            auc_scores[col] = 0.5
    return auc_scores


def _select_features_by_auc(X, y, threshold=0.51, plot=True, output_dir="plots"):
    """
    【改进】降低阈值至 0.51，避免过度剔除
    同时保留所有原始数值列（即使 AUC<0.51）
    新增：可视化 AUC 排名
    """
    selected = []
    original_numeric = [col for col in X.columns if X[col].dtype != 'object']

    auc_scores = _compute_feature_auc_scores(X, y)

    for col in X.columns:
        score = auc_scores[col]
        if score >= threshold:
            selected.append(col)
        elif col in original_numeric:
            selected.append(col)

    selected = list(dict.fromkeys(selected))

    # === 可视化 ===
    if plot:
        os.makedirs(output_dir, exist_ok=True)
        sorted_features = sorted(auc_scores.items(), key=lambda x: x[1], reverse=True)
        features, scores = zip(*sorted_features)

        plt.figure(figsize=(10, max(6, len(features) * 0.25)))
        colors = ['green' if f in selected else 'red' for f in features]
        plt.barh(range(len(features)), scores, color=colors, edgecolor='black')
        plt.yticks(range(len(features)), features)
        plt.xlabel("单变量 AUC (或 1-AUC)")
        plt.title("特征单变量区分能力 (绿色=保留, 红色=剔除)")
        plt.axvline(x=threshold, color='blue', linestyle='--', label=f'阈值={threshold}')
        plt.legend()
        plt.tight_layout()

        plot_path = os.path.join(output_dir, "feature_auc_ranking.png")
        plt.savefig(plot_path, dpi=150)
        plt.close()

    return X[selected]


def preprocess_for_training(y, X_raw):
    """
    【兼容】端到端训练数据预处理函数
    """
    X_eng = engineer_features(X_raw, is_train=True)
    X_final = _select_features_by_auc(X_eng, y, threshold=0.51, plot=True)
    return X_final, y


def preprocess_for_inference(X_raw):
    """
    【新增】用于测试集的预处理（与训练完全一致，但无筛选）
    """
    X_eng = engineer_features(X_raw, is_train=False)
    return X_eng
